import { memo, useCallback, useMemo, ChangeEvent } from 'react';
import { Select, Text } from '@chakra-ui/react';
import { useAppDispatch, useAppSelector } from 'store/storeHooks';

import { Flex } from '@chakra-ui/react';
import { createSelector } from '@reduxjs/toolkit';
import { stateSelector } from 'store/store';
import { defaultSelectorOptions } from 'store/util/defaultMemoizeOptions';
import { modelSelected } from 'features/parameters/store/actions';
import { MODEL_TYPE_MAP } from 'features/parameters/types/constants';
import { modelIdToMainModelParam } from 'features/parameters/util/modelIdToMainModelParam';
import { forEach } from 'lodash-es';
import { NON_REFINER_BASE_MODELS } from 'services/api/constants';
import { useGetMainModelsQuery } from 'services/api/endpoints/models';

const selector = createSelector(
  stateSelector,
  (state) => ({ model: state.generation.model }),
  defaultSelectorOptions
);

const ParamMainModelSelect = () => {
  const dispatch = useAppDispatch();

  const { model } = useAppSelector(selector);

  const { data: mainModels, isLoading } = useGetMainModelsQuery(
    NON_REFINER_BASE_MODELS
  );

  const data = useMemo(() => {
    if (!mainModels) {
      return [];
    }

    const data: any[] = [];

    forEach(mainModels.entities, (model, id) => {
      if (!model) {
        return;
      }

      data.push({
        value: id,
        label: model.model_name,
        group: MODEL_TYPE_MAP[model.base_model],
      });
    });

    return data;
  }, [mainModels]);

  // grab the full model entity from the RTK Query cache
  // TODO: maybe we should just store the full model entity in state?
  const selectedModel = useMemo(
    () =>
      mainModels?.entities[`${model?.base_model}/main/${model?.model_name}`] ??
      null,
    [mainModels?.entities, model]
  );

  const handleChangeModel = useCallback(
    (e: ChangeEvent) => {
      const v = (e.target as HTMLSelectElement).value;
      const newModel = modelIdToMainModelParam(v);

      if (!newModel) {
        return;
      }

      dispatch(modelSelected(newModel));
    },
    [dispatch]
  );

  return isLoading ? (
    <Text>Loading models...</Text>
  ) : (
    <Flex w="100%" alignItems="center" gap={3}>
      <Select
        value={selectedModel?.id}
        onChange={handleChangeModel}
        placeholder={data.length > 0 ? 'Select a model' : 'No models available'}
      >
        {data.map((m) => {
          return (
            <option key={m.value} value={m.value}>
              {m.label}
            </option>
          );
        })}
      </Select>
    </Flex>
  );
};

export default memo(ParamMainModelSelect);
